import datetime
import traceback
import argparse
import os
from Grade_5level import ResnetGrade
import time


grade_infer = ResnetGrade("./model-MESSIDOR.pth")

def run_one(file_path: str):
    try:
        t1 = time.perf_counter()
        out = grade_infer.infer(file_path)
        t2 = time.perf_counter()
        print("Infer file: {}, cost: {}s".format(os.path.basename(file_path), (t2 - t1)))
        return out
    except Exception as e:
        print(traceback.format_exc())
        return {
            'index': 0,
            'pie': [-1]
        }


def loop(from_path: str, csv_path: str):
    file_names = os.listdir(from_path)
    with open(csv_path, "w") as csv_file:
        csv_file.write("name,index,possible\n")
        for each_name in file_names:
            abs_path = os.path.join(from_path, each_name)
            out = run_one(abs_path)
            index = out["index"]
            possible = out["pie"]
            csv_file.write("{},{},{}\n".format(each_name, index, possible))


def run(work_dir: str):
    from_path = os.path.join(work_dir, "user_input")
    csv_name = "{}.csv".format(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S'))
    csv_path = os.path.join(work_dir, csv_name)
    if not os.path.exists(from_path):
        print("user_input dir is not found in {}, program will exit".format(work_dir))
        return

    print("Generate csv name is: {}".format(csv_name))

    loop(from_path, csv_path)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--workdir", type=str)
    args = parser.parse_args()
    if args.workdir == "" or args.workdir is None:
        print("work dir is empty")
        exit(-1)

    run(args.workdir)
